import datetime
import os
import threading
import time
import tkinter as tk
import tkinter.messagebox as tkm

import cv2
import logging
import numpy as np
from tkinterdnd2 import *
import tkinter.filedialog as tkf
from PIL import Image
import configparser
from time import strftime
from time import gmtime

from utils.utils_gui import imgConvert, set_logging, MyThread, LoggerBox, get_time, console_log, get_camrea, \
    gen_name, write_sql, get_uuid
from yolo import YOLO

content = configparser.ConfigParser()
content.read("config.ini", encoding='UTF-8')

# 选择的摄像头编号
camera_No = "0"
# 窗体尺寸
window_size = content.get("detect", "window_size")
# 网络权重
weight = content.get("detect", "weight")
# 置信度阈值
conf_thres = eval(content.get("detect", "conf_thres"))
# iou阈值
iou_thres = eval(content.get("detect", "iou_thres"))
# 背景图
bg = content.get("detect", "bg_image")
# cuda加速
cuda = content.get("train", "cuda")
# 窗口左上角图标
mask_icon = content.get("detect", "mask_icon")
# 画布
canvas = None
# 全局变量引用canvas，防止不生效
file_temp = None
# 输出文件夹
output_dir = content.get("detect", "output_dir")
# 主窗口
root = None
# 停止检测
stop_detect = False
# 日志
log_box, log_path = None, content.get("detect", "log_path")
# 重置系统，检测过程不可用
reset_system_btn = None
# 保存视频FPS
video_fps = eval(content.get("detect", "video_fps"))
# 摄像头信息和数目
camera_param, camera_count = None, 0
# 加载模型
model = None

################################################################# 控件状态更改实现 #############################################################
def change_widget_state(state=True):
    # True为启用控件，用于检测结束
    # False为禁用控件，用于检测开始

    # 重置系统控件
    load_reset_system_widget(state=state)
    # 停止检测控件（特例）
    load_stop_detect_widget(state=not state)
    # 选择权重控件
    load_select_weight_widget(state=state)
    # 开始检测控件
    load_start_detect_widget(state=state)
    # cuda加速控件
    load_cuda_acc_widget(state=state)
    # 写入sql空间
    load_write_sql_widget(state=state)



############################################################### 控件状态更改实现 END ###########################################################


################################################################# 检测实现 #############################################################
def start_detect():
    global conf_thres, iou_thres, output_dir, canvas, log_box, stop_detect, camera_No, model, video_fps, cuda

    if model == None:
        console_log(log_box, "检测终止，模型尚未初始化，请先加载模型")
        return
    # 开始检测，禁用部分控件
    change_widget_state(state=False)
    console_log(log_box, f"启用cuda加速！") if cuda else console_log(log_box, f"关闭cuda加速！")
    console_log(log_box, f"加载权重：{weight}")
    console_log(log_box, f"输出路径：{output_dir}")

    # 初始化改变停止标志防止无法检测
    stop_detect = False

    console_log(log_box, "即将进行摄像头检测！")
    console_log(log_box, f"使用{camera_No}号摄像头检测")

    capture = cv2.VideoCapture(int(camera_No))

    size = (int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    video_full_path = output_dir + gen_name(suffix=".mp4")
    # out = cv2.VideoWriter(video_full_path, cv2.VideoWriter_fourcc(*'H264'), video_fps, size)
    out = cv2.VideoWriter(video_full_path, cv2.VideoWriter_fourcc(*'avc1'), video_fps, size)

    ref, frame = capture.read()
    video_name = gen_name(suffix=".mp4")
    video_time = 0
    without_mask_flag = 0
    with_mask_flag = 0
    without_mask_count = 0
    with_mask_count = 0
    dangerous_level = ""
    fps = 0.0
    # noinspection PyRedundantParentheses
    while True and not stop_detect:
        t1 = time.time()
        # 读取某一帧
        ref, frame = capture.read()
        if not ref:
            break
        # 格式转变，BGRtoRGB
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        # 转变成Image
        frame = Image.fromarray(np.uint8(frame))
        # 推理
        image, classes_nums = model.camera_detect(log_box, frame, conf_thres, iou_thres)
        frame = np.array(image)
        # RGBtoBGR满足opencv显示格式
        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        fps = (fps + (1. / (time.time() - t1))) / 2
        # noinspection PyRedundantParentheses
        # print("fps= %.2f" % (fps))
        # noinspection PyRedundantParentheses
        frame = cv2.putText(frame, "fps= %.2f" % (fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

        # cv2.imshow("video", frame)
        # 绘图至canvas
        tkimg = imgConvert.convert(frame)
        canvas.create_image(0, 0, anchor='nw', image=tkimg, tag='im0')
        file_temp = tkimg
        # root.update()
        # root.after(1)
        # cv2.waitKey(1)
        out.write(frame)
        # 视频帧数+1
        video_time += 1

        # classes_nums[1, 1] classes_nums[without_mask, with_mask]
        # 1是检测到，0是未检测到
        without_mask_count += classes_nums[0]
        with_mask_count += classes_nums[1]

        # 停止检测
        if stop_detect:
            # 以tag为标记删除背景
            reset_canvas()
            console_log(log_box, "停止检测")
            # 停止检测标志更新
            stop_detect = False
            # 释放摄像头
            capture.release()
            # 保存视频流释放
            out.release()
            cv2.destroyAllWindows()
            # 窗体更新
            # root.update()
            # 检测结束，启用部分控件
            change_widget_state(state=True)

            console_log(log_box, "检测视频保存已保存至：" + output_dir + gen_name(suffix=".mp4"))

            without_mask_flag = 1 if without_mask_count else 0
            with_mask_flag = 1 if with_mask_count else 0
            # 未检测到对象
            if without_mask_count == 0 and with_mask_count == 0:
                dangerous_level = "安全"
            # 只有戴口罩的
            if without_mask_count == 0 and with_mask_count > 0:
                dangerous_level = "安全"
            # 只有未戴口罩的
            if without_mask_count > 0 and with_mask_count == 0:
                dangerous_level = "警告"
            # 戴口罩和未戴口罩的都有
            if without_mask_count > 0 and with_mask_count > 0:
                dangerous_level = "危险"

            video_time = strftime("%H:%M:%S", gmtime(video_time / video_fps))
            # 打开写入sql执行写入
            if write_sql_var.get():
                dt = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                threading.Thread(target=write_sql,
                                 args=(get_uuid(),
                                       video_name,
                                       video_name,
                                       video_time,
                                       without_mask_flag,
                                       with_mask_flag,
                                       dangerous_level,
                                       dt,
                                       dt,)).start()

############################################################### 检测实现 END ############################################################


############################################################### 初始化窗口参数 ###########################################################
def init_window():
    global root, camera_count, camera_param
    root = TkinterDnD.Tk()
    root.title('口罩佩戴检测')  # 窗口标题

    root.geometry(window_size)  # 窗口大小
    root.resizable(False, False)  # 禁止缩放
    root.config(bg='#F0F0F0')  # 窗口背景颜色
    root.iconbitmap(mask_icon)

    camera_param, camera_count = get_camrea()


init_window()


############################################################## 初始化窗口参数 END #########################################################


################################################################# 初始化画布 #############################################################
def load_canvas():
    global canvas
    canvas = tk.Canvas(root, bg='black', width=640, height=480)
    canvas.place(x=150, y=60)


load_canvas()


############################################################### 初始化画布 END ###########################################################


################################################################## 重置画布 ##############################################################
def reset_canvas():
    global canvas, file_temp, bg
    bg_ = cv2.imread(bg)
    bg_ = imgConvert.convert(bg_, True)
    file_temp = bg_
    canvas.create_image(0, 0, anchor='nw', image=bg_, tag='bg')


################################################################ 重置画布 END ###########################################################


############################################################### 初始化日志框 ###########################################################
def load_log_box():
    global log_box, root
    streamHandlerBox = LoggerBox(root, width=160, height=14)
    streamHandlerBox.place(x=80, y=570)
    log_box = logging.getLogger('log')
    log_box.setLevel(logging.INFO)
    log_box.addHandler(logging.StreamHandler(streamHandlerBox))


load_log_box()


############################################################## 初始化日志框 END ##########################################################


############################################################## 加载选择摄像头控件 ############################################################
def load_camera_ra_widget(state=True):
    global camera_count, camera_No
    camera_var = tk.IntVar()
    # 设备标签
    device_label = tk.Label(root, width=13, height=1, text='empty', font=("微软雅黑", 9, "bold"))
    device_label.place(x=860, y=70)

    # camera_var.set(0)
    device_label.config(text='摄像头「' + str(camera_var.get()) + '」')

    # 选择摄像头触发函数
    def camera_trigger():
        global camera_No
        device_label.config(text='摄像头「' + str(camera_var.get()) + '」')
        camera_No = str(camera_var.get())
        console_log(log_box, f"选择了{camera_No}号摄像头")

    for i in range(camera_count):
        # camera Radiobutton
        locals()["camera" + str(i + 1)] = tk.Radiobutton(root, text="# " + str(i), variable=camera_var, value=i,
                                                         command=camera_trigger, font=("微软雅黑", 9, "bold"),
                                                         cursor="hand2")
        locals()["camera" + str(i + 1)].place(x=1030 + i * 90, y=70)

    if camera_count == 0:
        no_device_label = tk.Label(root, width=13, height=1, text='无可用摄像头', font=("微软雅黑", 9, "bold"))
        no_device_label.place(x=1060, y=70)
        device_label.config(text='摄像头「无」')

    if not state:
        for i in range(camera_count):
            locals()["camera" + str(i + 1)].config(state='disable')
        console_log(log_box, "选择摄像头控件已禁用！")
    else:
        console_log(log_box, "选择摄像头控件加载成功！")


load_camera_ra_widget() if camera_count else load_camera_ra_widget(state=False)


############################################################# 加载选择摄像头控件 END #######################################################


############################################################### conf阈值控件 ##############################################################
def conf_thres_widget():
    conf_thres_label = tk.Label(root, text='conf-thres「' + str(conf_thres) + '」', font=("微软雅黑", 9, "bold"))
    conf_thres_label.place(x=840, y=120)

    def print_conf_thres_selection(value):
        global conf_thres
        conf_thres = float(value)
        conf_thres_label.config(text='conf-thres「' + str(conf_thres) + '」')

    # iou-thres Scale
    conf_thres_scale = tk.Scale(root, from_=0, to=1, orient=tk.HORIZONTAL, length=250, showvalue=0, tickinterval=0.2,
                                resolution=0.01,
                                command=print_conf_thres_selection, cursor="hand2")
    conf_thres_scale.place(x=1000, y=120)
    conf_thres_scale.set(conf_thres)
    console_log(log_box, "conf阈值控件加载成功！")


conf_thres_widget()


############################################################# conf阈值控件 END ############################################################


############################################################### iou阈值控件 ##############################################################
def iou_thres_widget():
    iou_thres_label = tk.Label(root, text='iou-thres「' + str(iou_thres) + '」', font=("微软雅黑", 9, "bold"))
    iou_thres_label.place(x=840, y=170)

    def print_iou_thres_selection(value):
        global iou_thres
        iou_thres = float(value)
        iou_thres_label.config(text='iou-thres「' + str(iou_thres) + '」')

    # iou-thres Scale
    iou_thres_scale = tk.Scale(root, from_=0, to=1, orient=tk.HORIZONTAL, length=250, showvalue=0, tickinterval=0.2,
                               resolution=0.01,
                               command=print_iou_thres_selection, cursor="hand2")
    iou_thres_scale.place(x=1000, y=170)
    iou_thres_scale.set(iou_thres)
    console_log(log_box, "iou阈值控件加载成功！")


iou_thres_widget()
############################################################ iou阈值控件 END ###########################################################


############################################################ 加载选择权重控件 ###########################################################
weight_path = tk.StringVar()


def load_select_weight_widget(state=True):
    def select_weight():
        global weight_path, weight, content, model
        weight_format = ['pt', 'pth']  # 可加载权重格式
        path = tkf.askopenfilename(filetypes=[("weight file", weight_format)])
        if path == weight:
            return
        # 模型加载完成后，权重加载不生效
        if path:
            weight_path.set(path)
            weight = path
            # 输出日志
            console_log(log_box, f"选择权重文件 {weight}")
            # 固定权重保存到配置文件中
            content.set("detect", "weight", path)
            content.write(open("config.ini", "w+", encoding="UTF-8"))
            # 根据模型是否加载输出权重更新成功与否
            if model != None:
                # tkm.showinfo(title="提示", message="模型加载后权重更新不成功，下次生效")
                console_log(log_box, "模型加载后权重更新不成功，下次生效")
            else:
                console_log(log_box, "权重更新成功")
                # tkm.showinfo(title="提示", message="权重更新成功")
        # 用户取消更改权重
        else:
            console_log(log_box, "用户取消更改权重")

    select_weight_entry = tk.Entry(root, textvariable=weight_path, width=45)
    select_weight_entry.place(x=850, y=240)

    weight_path.set(weight)

    select_weight_btn = tk.Button(root, text="选择权重", command=select_weight, relief=tk.RIDGE, bd=2,
                                  font=("微软雅黑", 9, "bold"), cursor="hand2")
    select_weight_btn.place(x=1175, y=236)

    if not state:
        select_weight_entry.config(state='disable')
        select_weight_btn.config(state='disable')
    console_log(log_box, "选择权重控件加载成功！")


load_select_weight_widget()

############################################################ 加载选择权重控件 END ###########################################################



############################################################ 加载视频输出文件夹选择控件 ###########################################################
output_dir_var = tk.StringVar()


# 初始化输出结果文件夹控件
def load_output_dir_widget(state=True):
    global output_dir

    def select_dir():
        global output_dir

        new_output_dir = tkf.askdirectory()
        if new_output_dir:
            output_dir_var.set(str(new_output_dir))
            output_dir = new_output_dir
            console_log(log_box, "视频输出文件夹已更改为："+new_output_dir)
    select_output_dir_entry = tk.Entry(root, textvariable=output_dir_var, width=45)
    select_output_dir_entry.place(x=850, y=290)

    output_dir_var.set(output_dir)

    select_output_dir_btn = tk.Button(root, text="输出位置", command=select_dir, relief=tk.RIDGE, bd=2,
                                      font=("微软雅黑", 9, "bold"), cursor="hand2")
    select_output_dir_btn.place(x=1175, y=286)

    if not state:
        select_output_dir_entry.config(state='disable')
        select_output_dir_btn.config(state='disable')
        console_log(log_box, "选择输出文件夹控件已禁用！")
    else:
        console_log(log_box, "选择输出文件夹控件加载成功！")


load_output_dir_widget(state=True)
############################################################ 加载视频输出文件夹选择控件 END ###########################################################


############################################################# 视频写入数据库控件 ############################################################
write_sql_var = tk.IntVar()


def load_write_sql_widget(state=True):
    global write_sql_var
    def write_sql():
        global write_sql_var
        if write_sql_var.get() == 1:

            console_log(log_box, "写数据库控件已启用！")
        else:

            console_log(log_box, "写数据库控件已禁用！")

    write_sql_dir_cbtn = tk.Checkbutton(root, text='写数据库', variable=write_sql_var, onvalue=1, offvalue=0,
                                        command=write_sql, font=("微软雅黑", 9, "bold"), cursor="hand2")
    write_sql_dir_cbtn.place(x=845, y=330)
    write_sql_var.set(1)

    if not state:
        write_sql_dir_cbtn.config(state='disable')
        console_log(log_box, "写数据库控件已禁用！")
    else:
        console_log(log_box, "写数据库控件加载成功！")


load_write_sql_widget(state=True)

########################################################### 视频写入数据库控件 END ##########################################################


############################################################# 使用cuda加速控件 ############################################################
cuda_acc_var = tk.IntVar()

def load_cuda_acc_widget(state=True):
    global cuda_acc_var
    def cuda_acc():
        global cuda_acc_var
        if cuda_acc_var.get() == 1:
            console_log(log_box, "cuda加速已启用！")
        else:
            console_log(log_box, "cuda加速已关闭！")

    cuda_acc_cbtn = tk.Checkbutton(root, text='cuda加速', variable=cuda_acc_var, onvalue=1, offvalue=0,
                                        command=cuda_acc, font=("微软雅黑", 9, "bold"), cursor="hand2")
    cuda_acc_cbtn.place(x=950, y=330)
    cuda_acc_var.set(1)

    if not state:
        cuda_acc_cbtn.config(state='disable')
        console_log(log_box, "cuda加速控件已禁用！")
    else:
        console_log(log_box, "cuda加速控件已启用！")


load_cuda_acc_widget(state=True)

########################################################### 视频写入数据库控件 END ##########################################################


############################################################### 预加载模型控件 ##############################################################
def load_model_widget():
    def load_model():
        global weight, model, cuda, camera_No

        model = YOLO(cuda=cuda)

        # 初始化日志
        set_logging()

        # tkm.showinfo(title="提示", message="模型加载成功")
        console_log(log_box, f"主干网络权重：{weight}")
        console_log(log_box, f"默认使用cuda加速") if cuda else console_log(log_box, f"不使用cuda加速")
        console_log(log_box, "模型预加载成功！")

        # 模型加载完，禁用cuda选择
        load_cuda_acc_widget(state=False)

    load_model_btn = tk.Button(root, text='加载模型', command=lambda: MyThread(load_model), relief=tk.RIDGE, bd=2,
                               font=("微软雅黑", 9, "bold"), cursor="hand2")
    load_model_btn.place(x=40, y=80)
    console_log(log_box, "加载模型控件加载成功！")


load_model_widget()


############################################################# 预加载模型控件 END #############################################################



############################################################## 加载开始检测控件 #############################################################
def load_start_detect_widget(state):
    # 打开摄像头按钮
    start_detect_btn = tk.Button(root, text='开始检测', command=lambda: MyThread(start_detect), relief=tk.RIDGE, bd=2,
                                font=("微软雅黑", 9, "bold"), cursor="hand2")
    start_detect_btn.place(x=40, y=160)

    if not state:
        start_detect_btn.config(state='disable')
        console_log(log_box, "停止检测控件已禁用！")
    else:
        console_log(log_box, "开始检测控件加载成功！")

load_start_detect_widget(state=True)


############################################################ 加载开始检测控件 END ###########################################################


############################################################# 加载停止检测控件 ############################################################
def load_stop_detect_widget(state=True):
    def stop_detect():
        global stop_detect
        stop_detect = True

    # 打开摄像头按钮
    stop_detect_btn = tk.Button(root, text='停止检测', command=lambda: MyThread(stop_detect), relief=tk.RIDGE, bd=2,
                                font=("微软雅黑", 9, "bold"), cursor="hand2")
    stop_detect_btn.place(x=40, y=240)

    if not state:
        stop_detect_btn.config(state='disable')
        console_log(log_box, "停止检测控件已禁用！")
    else:
        console_log(log_box, "停止检测控件加载成功！")


load_stop_detect_widget(state=False)


############################################################ 加载停止检测控件 END ##########################################################


############################################################# 加载重置系统控件 ############################################################
def load_reset_system_widget(state=True):
    global reset_system_btn

    def reset_system():
        global canvas, model
        if model != None:
            model = None
            # 模型置空，设备选择释放，且
            load_cuda_acc_widget(state=True)

        reset_canvas()

        # tkm.showinfo(title="提示", message="系统重置成功")
        console_log(log_box, "系统重置成功")

    reset_system_btn = tk.Button(root, text='重置系统', command=lambda: MyThread(reset_system), relief=tk.RIDGE, bd=2,
                                 font=("微软雅黑", 9, "bold"), cursor="hand2")
    reset_system_btn.place(x=40, y=320)

    if not state:
        reset_system_btn.config(state='disable')
        console_log(log_box, "重置系统控件已禁用！")
    else:
        console_log(log_box, "重置系统控件加载成功！")


load_reset_system_widget()


############################################################ 加载重置系统控件 END ##########################################################


############################################################## 加载检测结果控件 ############################################################
def load_detect_result_widget():
    def open_result():
        global output_dir

        if output_dir == None:
            console_log(log_box, "暂无运行结果，请检测后再打开")
        else:
            output_dir = str(output_dir)
            # 找不到:说明路径为相对路径 runs\detect\exp4
            if output_dir.find(":") == -1:
                output_dir = os.getcwd() + "\\" + output_dir
            try:
                os.startfile(output_dir)
            except FileNotFoundError:
                console_log(log_box, "暂无结果，请检测后再打开")

    detect_result_btn = tk.Button(root, text='检测结果', command=open_result, relief=tk.RIDGE, bd=2,
                                  font=("微软雅黑", 9, "bold"), cursor="hand2")
    detect_result_btn.place(x=40, y=400)
    console_log(log_box, "检测结果控件加载成功！")


load_detect_result_widget()


############################################################ 加载检测结果控件 END ##########################################################


################################################################ 加载获取设备按钮 ############################################################
def load_get_device_widget():
    def get_camrea_info():
        global camera_param, camera_count
        console_log(log_box, "▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼设备信息▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼")
        console_log(log_box, camera_param)
        console_log(log_box, f"该设备共有{camera_count}个摄像头")

    get_device_btn = tk.Button(root, text='获取设备', command=get_camrea_info, relief=tk.RIDGE, bd=2,
                               font=("微软雅黑", 9, "bold"), cursor="hand2")
    get_device_btn.place(x=40, y=480)


load_get_device_widget()


############################################################# 加载获取设备按钮 END #########################################################


################################################################# 手动关闭程序 ############################################################
def close():
    if tkm.askokcancel("关闭", "是否关闭程序"):
        console_log(log_box, "程序退出")
        cv2.destroyAllWindows()
        root.destroy()


############################################################### 手动关闭程序 END ##########################################################

console_log(log_box, "控件全部加载成功！")
console_log(log_box, "▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼当前信息▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼")
console_log(log_box, f"摄像头：{camera_No}号")
console_log(log_box, f"权重文件：{weight}")
console_log(log_box, f"检测结果目录：{output_dir}")
console_log(log_box, f"启用cuda加速！") if cuda else console_log(log_box, f"关闭cuda加速！")
console_log(log_box, f"视频开启写入SQL！") if write_sql_var else console_log(log_box, f"关闭cuda加速！")


root.protocol("WM_DELETE_WINDOW", close)
root.mainloop()
